import torch
import matplotlib.pyplot as plt

x1 = torch.normal(1, 1, (20,))
x2 = torch.normal(1, 1, (20,))
x1 = torch.sort(x1).values
x2 = torch.sort(x2).values

# noise1 = torch.normal(0, 0.3, (20,))
# noise2 = torch.normal(0, 0.3, (20,))

# ------------------多特征输入------------------------
fig, ax = plt.subplots(subplot_kw={"projection": "3d"})

y = 2 * x1 + 3 * x2 + 3

ax.set_xlabel("x1")
ax.set_ylabel("x2")
ax.set_zlabel("y")
ax.plot(x1.detach().numpy(), x2.detach().numpy(), y.detach().numpy(), 'r-')

w1 = torch.tensor(0.1, requires_grad=True)
w2 = torch.tensor(0.2, requires_grad=True)
b = torch.tensor(0.3, requires_grad=True)

epochs = 100
for epoch in range(epochs):
    y_predict = w1 * x1 + w2 * x2 + b
    e = torch.mean((y_predict - y) ** 2)
    e.backward()
    with torch.no_grad():
        w1 -= w1.grad * 0.01
        w2 -= w2.grad * 0.01
        b -= b.grad * 0.01
        w1.grad.zero_()
        w2.grad.zero_()
        b.grad.zero_()

    print(e)
y_predict = w1 * x1 + w2 * x2 + b
ax.plot(x1.detach().numpy(), x2.detach().numpy(), y_predict.detach().numpy(), 'b-')
plt.show()
